import uuid
import time
import redis


class RedisLock:
    LUA_RELEASE = """
        if redis.call("GET", KEYS[1]) == ARGV[1] then
            return redis.call("DEL", KEYS[1])
        else
            return 0
        end
    """

    LUA_EXTEND = """
        if redis.call("GET", KEYS[1]) == ARGV[1] then
            return redis.call("PEXPIRE", KEYS[1], ARGV[2])
        else
            return 0
        end
    """

    def __init__(self, redisInstance: redis.Redis, key: str, appId: str,
                 expire: int = 5,  # seconds
                 retry: int = 15,  # max retry times
                 retry_delay: float = 0.2):
        self.r = redisInstance
        self.key = f"lock:{key}:{appId}"
        self.expire_ms = int(expire * 1000)
        self.retry = retry
        self.retry_delay = retry_delay
        self.token = None

    def acquire(self, appId: str) -> bool:
        self.token = uuid.uuid4().hex
        for _ in range(self.retry):
            if self.r.set(self.key + appId, self.token, nx=True, px=self.expire_ms):
                return True
            time.sleep(self.retry_delay)
        return False

    def release(self, appId: str) -> bool:
        if self.token is None:
            return False
        return bool(self.r.eval(self.LUA_RELEASE, 1, self.key + appId, self.token))

    def extend(self, appId: str, extra_sec: int = 5) -> bool:
        if self.token is None:
            return False
        return bool(self.r.eval(self.LUA_EXTEND, 1, self.key + appId, self.token, int(extra_sec * 1000)))

